from tensorflow.keras.datasets import mnist
from tensorflow.keras.utils import to_categorical

from tensorflow.keras import layers
from tensorflow.keras.models import Model,load_model
from tensorflow.keras import optimizers
from tensorflow.keras.losses import categorical_crossentropy






(x_train,y_train),(x_test,y_test) = mnist.load_data()
X_train = x_train.reshape((60000,28,28,1))/255.
X_test = x_test.reshape((10000,28,28,1))/255.

y_train = to_categorical(y_train,num_classes=10)
y_test = to_categorical(y_test,num_classes=10)

temperature = 3
class Calculate_loss(layers.Layer):
    def __init__(self,Temperature,label_loss_weight,**kwargs):
        self.temperature = Temperature
        self.label_loss_weight = label_loss_weight
        super(Calculate_loss,self).__init__(**kwargs)

    def call(self,inputs):
        student_output = inputs[0]
        teacher_output = inputs[1]
        labels = inputs[2]
        loss_1 = categorical_crossentropy(teacher_output/self.temperature,student_output/self.temperature)
        loss_2 = self.label_loss_weight*categorical_crossentropy(labels,student_output)
        self.add_loss(loss_1 + loss_2,inputs=inputs)
        return student_output
teacher = load_model('teacher_model/')
teacher.trainable = False

y_inputs = layers.Input((10,))
y = layers.Lambda(lambda t:t)(y_inputs)
inputs = layers.Input((28,28,1))
x = layers.Conv2D(16,3,activation='relu')(inputs)
x = layers.Conv2D(16,3,strides=2,activation='relu')(x)
x = layers.Conv2D(32,5,activation='relu')(x)
x = layers.Conv2D(32,5,activation='relu')(x)
x = layers.Flatten()(x)
x = layers.Dense(60,activation='relu')(x)
x = layers.Dense(10,activation='softmax')(x)
x = Calculate_loss(Temperature=temperature,label_loss_weight=0.1)([x,teacher(inputs),y])
x = layers.Lambda(lambda t:t/temperature)(x)
student = Model([inputs,y_inputs],x)
student.compile(optimizer=optimizers.SGD(momentum=0.9,nesterov=True),loss=None)
student.summary()
student.fit(x=[X_train,y_train],y=None,batch_size=128,epochs=1,verbose=2)
student.save("student_with_hard_label/")
softmax_layer=student.layers[-4]

predict_model=Model(inputs,softmax_layer.output)

res=predict_model.predict(X_test)

import numpy as np
result=[np.argmax(a) for a in res]

(x_train,y_train),(x_test,y_test)=mnist.load_data()

from sklearn.metrics import accuracy_score
accuracy_score(y_test,result)# 98.79%
